import logging, os, time

from .util import generate_data_path
from pathlib import Path
from playwright.async_api import async_playwright
from jinja2 import Template
from typing_extensions import TypedDict
from playwright.async_api._generated import Locator
from playwright._impl._api_structures import FloatRect
from typing import Literal

class ScreenshotOptions(TypedDict):
    """Playwright 截图参数

    详见：https://playwright.dev/python/docs/api/class-page#page-screenshot

    Args:
        timeout (float, optional): 截图超时时间.
        type (Literal["jpeg", "png"], optional): 截图图片类型.
        path (Union[str, Path]], optional): 截图保存路径，如不需要则留空.
        quality (int, optional): 截图质量，仅适用于 JPEG 格式图片.
        omit_background (bool, optional): 是否允许隐藏默认的白色背景，这样就可以截透明图了，仅适用于 PNG 格式.
        full_page (bool, optional): 是否截整个页面而不是仅设置的视口大小，默认为 True.
        clip (FloatRect, optional): 截图后裁切的区域，xy为起点.
        animations: (Literal["allow", "disabled"], optional): 是否允许播放 CSS 动画.
        caret: (Literal["hide", "initial"], optional): 当设置为 `hide` 时，截图时将隐藏文本插入符号，默认为 `hide`.
        scale: (Literal["css", "device"], optional): 页面缩放设置.
            当设置为 `css` 时，则将设备分辨率与 CSS 中的像素一一对应，在高分屏上会使得截图变小.
            当设置为 `device` 时，则根据设备的屏幕缩放设置或当前 Playwright 的 Page/Context 中的
            device_scale_factor 参数来缩放.
        mask (List["Locator"]], optional): 指定截图时的遮罩的 Locator。元素将被一颜色为 #FF00FF 的框覆盖.

    @author: Redlnn(https://github.com/GraiaCommunity/graiax-text2img-playwright)
    """

    timeout: float | None
    type: Literal["jpeg", "png", None]
    quality: int | None
    omit_background: bool | None
    full_page: bool | None
    clip: FloatRect | None
    animations: Literal["allow", "disabled", None]
    caret: Literal["hide", "initial", None]
    scale: Literal["css", "device", None]
    mask: list[Locator] | None

class Text2ImgRender:
    def __init__(self):
        self.playwright = None
        self.browser = None
        self.context: BrowserContext = None

    async def from_jinja_template(self, template: str, data: dict) -> str:
        html = Template(template).render(data)
        return await self.from_html(html)
    
    async def from_html(self, html: str) -> str:
        html_file_path, abs_path = generate_data_path(suffix="html", namespace="rendered")
        with open(html_file_path, "w", encoding='utf-8') as f:
            f.write(html)
        return html_file_path, abs_path
    
    async def html2pic(self, html_file_path: str, screenshot_options: ScreenshotOptions) -> str:
        if self.context is None:
            self.playwright = await async_playwright().start()
            # 启动浏览器时增加 args 参数，以禁用一些可能影响截图质量的特性
            # headless=True 保持无头模式，方便服务器部署
            self.browser = await self.playwright.chromium.launch(
                headless=True,
                args=[
                    '--disable-gpu',
                    '--no-sandbox',
                    '--disable-setuid-sandbox',
                    '--disable-dev-shm-usage',
                    '--force-device-scale-factor=2', # 关键：强制设备像素比，提高清晰度
                    '--high-dpi-support=1' # 关键：启用高DPI支持
                ]
            )
            # 增加 device_scale_factor 到 context 中，进一步提高分辨率
            # 默认的viewport也可以根据实际需求调整，例如为了适应移动设备尺寸
            self.context = await self.browser.new_context(
                viewport={"width": 390, "height": 844},
                device_scale_factor=2 # 关键：设备像素比，通常2可以提供Retina显示效果
            )

        result_path, _ = generate_data_path(suffix="jpeg", namespace="rendered")

        logging.info(f"Rendering {html_file_path}")
        
        page = await self.context.new_page()
        await page.goto(f'file://{html_file_path}')
        
        # 尝试获取内容容器
        container = await page.query_selector('[data-screenshot-container]')
        
        # 移除 full_page 选项，因为它与元素截图冲突
        if 'full_page' in screenshot_options:
            screenshot_options.pop('full_page')
        
        # 增加 quality 参数以提高 JPG 压缩质量，默认为80，可以尝试90或100
        # 也可以指定 type='png' 获取无损图片，但文件大小会更大
        default_screenshot_options = {
            'type': 'jpeg',
            'quality': 90, # 提高JPEG质量
            'scale': 'device', # 确保使用设备的像素比
        }
        # 合并传入的选项，传入的选项会覆盖默认值
        final_screenshot_options = {**default_screenshot_options, **screenshot_options}

        if container:
            await container.screenshot(path=result_path, **final_screenshot_options)
        else:
            logging.warning("No element with data-screenshot-container found. Taking full page screenshot.")
            # 如果没有找到特定容器，则截取整个页面，此时可以考虑full_page=True
            await page.screenshot(path=result_path, full_page=True, **final_screenshot_options)
            
        await page.close()

        logging.info(f"Rendered {html_file_path} to {result_path}")
        return result_path

    async def close(self):
        """添加一个关闭浏览器上下文和Playwright实例的方法"""
        if self.context:
            await self.context.close()
            self.context = None
        if self.browser:
            await self.browser.close()
            self.browser = None
        if self.playwright:
            await self.playwright.stop()
            self.playwright = None
